package cluster

import (
	"strings"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
)

type K8sResource struct {
	MilliCPU float64
	Memory   float64
	// ScalarResources
	ScalarResources map[v1.ResourceName]float64
}

// Add is used to add two given resources
func (r *K8sResource) Add(rr *K8sResource) *K8sResource {
	r.MilliCPU += rr.MilliCPU
	r.Memory += rr.Memory

	for rName, rQuant := range rr.ScalarResources {
		if r.ScalarResources == nil {
			r.ScalarResources = map[v1.ResourceName]float64{}
		}
		r.ScalarResources[rName] += rQuant
	}
	return r
}

func (r *K8sResource) AddGPU(rr *K8sResource) *K8sResource {
	for rName, rQuant := range rr.ScalarResources {
		if r.ScalarResources == nil {
			r.ScalarResources = map[v1.ResourceName]float64{}
		}
		r.ScalarResources[rName] += rQuant
	}

	return r
}

// SetScalar sets a resource by a scalar value of this resource.
func (r *K8sResource) SetScalar(name v1.ResourceName, quantity float64) {
	// Lazily allocate scalar resource map.
	if r.ScalarResources == nil {
		r.ScalarResources = map[v1.ResourceName]float64{}
	}
	r.ScalarResources[name] = quantity
}

// AddScalar adds a resource by a scalar value of this resource.
func (r *K8sResource) AddScalar(name v1.ResourceName, quantity float64) {
	r.SetScalar(name, r.ScalarResources[name]+quantity)
}

func (r *K8sResource) AddPodRequest(podSpec *v1.PodSpec) *K8sResource {
	for _, container := range podSpec.Containers {
		if IsGPURequiredContainer(&container) {
			r.AddGPU(NewK8sResource(container.Resources.Requests))
		} else {
			r.Add(NewK8sResource(container.Resources.Requests))
		}
	}
	return r
}

func (r *K8sResource) buildK8sAllocatedResourceList() v1.ResourceList {
	result := v1.ResourceList{}
	if r.MilliCPU > 0 {
		result[v1.ResourceCPU] = *resource.NewMilliQuantity(int64(r.MilliCPU), resource.DecimalSI)
	}
	if r.Memory > 0 {
		result[v1.ResourceMemory] = *resource.NewQuantity(int64(r.Memory), resource.BinarySI)
	}

	for rName, rQuant := range r.ScalarResources {
		if rQuant > 0 {
			if strings.HasPrefix(string(rName), v1.ResourceHugePagesPrefix) {
				result[rName] = *resource.NewQuantity(int64(rQuant), resource.BinarySI)
			} else {
				result[rName] = *resource.NewQuantity(int64(rQuant), resource.DecimalSI)
			}
		}
	}
	if len(result) == 0 {
		return nil
	}
	return result
}

type Resource struct {
	CPU     int    `json:"cpu"`
	Memory  int64  `json:"memory"`
	GpuType string `json:"gpu_type,omitempty"`
	GpuNum  int    `json:"gpu_num,omitempty"`
}

type NodeInfo struct {
	Cpu     int64
	Mem     int64
	GpuType string // gpu类型,对应gpu-k8s-resource-name
	GpuNum  int64  // gpu数量
}

type NodeAvailResourceSummary struct {
	Name         string   `json:"name"`
	Role         string   `json:"role"`
	CpuModelName string   `json:"cpu_model_name"`
	WorkerType   string   `json:"worker_type"`
	IsTaint      bool     `json:"taint"`
	Available    Resource `json:"available"`
}
